#include "vm.h"
#include "pmm.h"
#include "kalloc.h"
#include "string.h"
#include "printf.h"
#include "riscv.h"

extern char etext[]; // 内核代码结束地址

pagetable_t kernel_pagetable;

// 构造SATP寄存器值
#define MAKE_SATP(pagetable) (0x8000000000000000 | ((uint64)pagetable >> 12))

pte_t* walk(pagetable_t pagetable, uint64 va, int alloc) {
    if(va >= MAXVA)
        return 0;
    
    for(int level = 2; level > 0; level--) {
        pte_t *pte = &pagetable[PX(level, va)];
        if(*pte & PTE_V) {
            pagetable = (pagetable_t)PTE2PA(*pte);
        } else {
            if(!alloc || (pagetable = (pagetable_t)alloc_page()) == 0)
                return 0;
            memset(pagetable, 0, PGSIZE);
            *pte = PA2PTE(pagetable) | PTE_V;
        }
    }
    return &pagetable[PX(0, va)];
}

int map_pages(pagetable_t pt, uint64 va, uint64 pa, uint64 size, int perm) {
    uint64 a, last;
    
    if(size == 0) return -1;
    
    a = PGROUNDDOWN(va);
    last = PGROUNDDOWN(va + size - 1);
    for(;;) {
        if(map_page(pt, a, pa, perm) < 0)
            return -1;
        if(a == last)
            break;
        a += PGSIZE;
        pa += PGSIZE;
    }
    return 0;
}

int map_page(pagetable_t pt, uint64 va, uint64 pa, int perm) {
    if(va % PGSIZE != 0 || pa % PGSIZE != 0) {
        printf("map_page: unaligned va=%p pa=%p\n", va, pa);
        return -1;
    }

    pte_t *pte = walk(pt, va, 1);  // alloc=1表示必要时创建
    if(!pte) {
        printf("map_page: walk failed\n");
        return -1;
    }
    
    if(*pte & PTE_V) {
        printf("map_page: already mapped va=%p\n", va);
        return -1;
    }

    *pte = PA2PTE(pa) | perm | PTE_V;
    return 0;
}
//内核页表初始化
void kvminit(void) {
printf("Initializing VM, etext at %p\n", etext);
    kernel_pagetable = (pagetable_t)alloc_page();
    uint64 align_etext = PGROUNDUP((uint64)etext);
    memset(kernel_pagetable, 0, PGSIZE);
    
    // 映射内核代码段 (RX)
    map_pages(kernel_pagetable, KERNBASE, KERNBASE, 
             align_etext - KERNBASE, PTE_R | PTE_X);
    
    // 映射内核数据段 (RW)
    map_pages(kernel_pagetable, (uint64)etext, (uint64)etext,
             PHYSTOP - (uint64)etext, PTE_R | PTE_W);
    
    // 映射设备
    map_pages(kernel_pagetable, UART0, UART0, PGSIZE, PTE_R | PTE_W);
}

void kvminithart(void) {
    w_satp(MAKE_SATP(kernel_pagetable));
    sfence_vma();
}

pagetable_t create_pagetable() {
void *page = kalloc();  // 使用void*中间转换
    if(!page) return NULL;
    memset(page, 0, PGSIZE);
    return (pagetable_t)page;
}
pte_t* walk_lookup(pagetable_t pt, uint64 va) {
    return walk(pt, va, 0);  // alloc=0表示仅查找
}

void free_pagetable(pagetable_t pt) {
    // 递归释放所有下级页表
    for(int i=0; i<512; i++) {
        if(pt[i] & PTE_V && !(pt[i] & (PTE_R|PTE_W|PTE_X))) {
            free_pagetable((pagetable_t)PTE2PA(pt[i]));
        }
    }
    kfree(pt);
}
